import os, shutil
import sys
import os.path as osp
import json
import numpy as np
import cv2
import copy
from multiprocessing import Pool
from tqdm import tqdm

# 递归获取文件夹下所有符合条件的文件路径.
def get_all_filepath(root_dir_path, extension_tag=[], recursive = True):
    """
    Args:
        root_dir_path: 文件根目录.
        extension_tag: 文件后缀名组成的列表, Example: ["jpg", "png", "bmp"]
        如果为空 则获取所有
    Returns:
        特定文件夹下所有符合条件的文件路径.
    """
    def check_endswith(new_path, extension_tag):
        if len(extension_tag) == 0:
            return True
        for s in extension_tag:
            if new_path.lower().endswith(s):
                return True
        return False

    paths = []
    for f in os.listdir(root_dir_path):
        new_path = root_dir_path + os.sep + f
        if os.path.isfile(new_path) and check_endswith(new_path,  extension_tag):
            paths.append(new_path)
        elif os.path.isdir(new_path) and recursive:
            temp_list = get_all_filepath(new_path, extension_tag)
            paths.extend(temp_list)
    return paths


def txt_to_via(data_dir, 
                via_name="via_region_data.json"):
    """
    ua detrac tool 
    将此数据集转换成via格式
    """

    via_files = get_all_filepath(data_dir, [via_name])

    thread_count  = min(12, len(via_files)) 
    p = Pool(thread_count)
    for f in via_files:
        p.apply_async(thread_deal, args=(f,))
        # thread_deal(f)
    # print('Waiting for all subprocesses done...')
    p.close()
    p.join()


def thread_deal(via_path):

    save_data = dict()
    data_dir = osp.dirname(via_path)
    with open(via_path,"r") as rf:
        data_dict = json.loads(rf.read())

    for key, data in tqdm(data_dict.items()):
        regions = []
        if len(data["regions"]) == 0: continue
        filename = data["filename"]
        img_path = osp.join(data_dir, filename)
        img = cv2.imread(img_path)

        rgm_mask = np.zeros(img.shape[:2], dtype=np.uint8)
        mask = np.zeros(img.shape[:2], dtype=np.uint8)
        for region in data["regions"]:
            rect = [
                region["shape_attributes"]["x"],
                region["shape_attributes"]["y"],
                region["shape_attributes"]["width"],
                region["shape_attributes"]["height"]
            ]
            if region["region_attributes"]["label"] == "ringelman":
                regions.append(region)
                cv2.rectangle(rgm_mask, (rect[0], rect[1]), (rect[0]+rect[2], rect[1]+rect[3]), [255], -1)
            else:
                cv2.rectangle(mask, (rect[0], rect[1]), (rect[0]+rect[2], rect[1]+rect[3]), [255], -1)
        if np.sum(rgm_mask) == 0: continue

        # cv2.imwrite("rgm_mask.jpg", rgm_mask)
        # cv2.imwrite("mask.jpg", mask)

        mask = mask - rgm_mask

        # cv2.imwrite("mask1.jpg", mask)
        mask = cv2.bitwise_not(mask)
        # cv2.imwrite("mask2.jpg", mask)

        bg = cv2.GaussianBlur(img, (51, 51), 80.)
        bg = cv2.copyTo(img, mask, bg)
        cv2.imwrite(img_path, bg)
        # cv2.imwrite("bg.jpg", bg)

  
        one_json = data
        one_json["regions"] = regions
        size = osp.getsize(img_path)
        one_json["size"] = size
        new_key = filename + str(size)
        save_data[new_key] = one_json

    with open(via_path, "w") as wf:
        wf.write(json.dumps(save_data))

def show_img_add_rect(data_dir, sub_name):
    """
    通过连续播放图片来增加过滤框
    用于过滤行人区域，避免标记
    """
    image_dir = osp.join(data_dir, "Insight-MVT_Annotation_Train")
    label_gt= osp.join(data_dir, "train_gt.txt")
    label_ign= osp.join(data_dir, "train_ign.txt")
    assert osp.exists(image_dir)
    assert osp.exists(label_gt)
    assert osp.exists(label_ign)

    # 每一张图的数据
    data_dict = dict()
    # 每一个文件夹的过滤数据
    ignore_dict = dict()
    with open(label_gt, "r") as rf:
        for line in rf.readlines():
            info = line.strip().split(" ")
            data_dict[info[0]] = np.array(info[1:], dtype=np.float32).reshape(-1, 4).astype(np.int)
    with open(label_ign, "r") as rf:
        for line in rf.readlines():
            info = line.strip().split(" ")
            ignore_dict[info[0]] = np.array(info[1:], dtype=np.float32).reshape(-1, 4).astype(np.int)
        

    sub_dir = osp.join(image_dir, sub_name)
    files = [[f, int(f[3:8])] for f in os.listdir(sub_dir)]
    files = sorted(files, key=lambda x: x[1])
    for f,idx in files:
        img = cv2.imread(osp.join(sub_dir, f))
        for ign_rect in ignore_dict.get(sub_name, []): 
            cv2.rectangle(img, (ign_rect[0], ign_rect[1]), (ign_rect[2], ign_rect[3]), [0,0,255], 1)
        for rect in data_dict.get(osp.join(sub_name, f), []):
            cv2.rectangle(img, (rect[0], rect[1]), (rect[2], rect[3]), [0,255,0], 1)
        cv2.imshow("image", img)
        cv2.waitKey(10)


def mask_by_label(data_dir, via_name="via_region_data.json"):
    via_files = get_all_filepath(data_dir, [via_name])

    
    for via_path in via_files:
        data_dir = osp.dirname(via_path)
        with open(via_path,"r") as rf:
            data_dict = json.loads(rf.read())

        for key, data in tqdm(data_dict.items()):
            if len(data["regions"]) == 0: continue
            filename = data["filename"]
            img_path = osp.join(data_dir, filename)
            print(img_path)

            for region in data["regions"]:

                if region["region_attributes"]["label"] != "mask":
                    continue
                
                points = np.array([[x,y] for x,y in zip(
                                region['shape_attributes']["all_points_x"], 
                                region['shape_attributes']["all_points_y"])], 
                                dtype=np.int)
                img = cv2.imread(img_path)
                bg = np.zeros(img.shape, dtype=np.uint8)
                mask = np.zeros(img.shape[:2], dtype=np.uint8)
                cv2.fillConvexPoly(mask, points, [255])
                mask = cv2.bitwise_not(mask, mask)
                bg = cv2.copyTo(img, mask, bg)
                # cv2.imwrite("1.jpg", bg)
                cv2.imwrite(img_path, bg)

                break

    
if __name__ == "__main__":
    # # 将数据集根据mask转换
    # data_dir = "/home/xc/work/code/paddle/train_data/det/ringelman/images/2021-06-29"
    # txt_to_via(data_dir)

    # 将 label = mask 的掩盖
    data_dir = sys.argv[1]
    # data_dir = "/home/xc/work/code/paddle/train_data/det/ringelman/images/2021-08-07"
    mask_by_label(data_dir)

